import time
import airsim
import numpy as np
import cvxpy as cp
import matplotlib.pyplot as plt
import control
from mpl_toolkits.mplot3d import Axes3D
import math



class Quadrotor:
    def __init__(self,dt):
        # 物理参数
        inertia = self.compute_inertia_matrix([0.18, 0.11, 0.04], [[0, 0.2275, 0.025],
                                                              [0, -0.2275, 0.025],
                                                              [0.2275, 0, 0.025],
                                                              [-0.2275, 0, 0.025]], 0.78, 0.055)
        self.mass = 1 # 质量 (kg)
        self.g = 9.8 # 重力加速度 (m/s²)
        self.Ixx =  inertia[0,0]  # 绕x轴转动惯量 (kg·m²)
        self.Iyy = inertia[1,1]   # 绕y轴转动惯量 (kg·m²)
        self.Izz = inertia[2,2]
        self.dt = dt  # 时间步长 (s)

        # 状态变量：[x, y, z, vx, vy, vz, roll, pitch,yaw, p, q, r]
        self.n_states = 12
        self.n_controls = 4  # [F,M1,M2,M3]

        # 状态空间模型
        self.A,self.B = self.get_dLTI()

    def compute_inertia_matrix(self,body_box, rotor_poses, box_mass, motor_assembly_weight):
        """
        计算无人机惯性矩阵的Python实现
        :param body_box: 机体尺寸数组[x,y,z] (单位：米)默认值[0.18,0.11,0.04]
        :param rotor_poses: 旋翼位置列表[[x1,y1,z1], [x2,y2,z2],...] (单位：米)
        [[0,0.2275,0.025],
       [0,-0.2275,0.025],
       [0.2275,0,0.025],
       [-0.2275,0,0.025]]
        :param box_mass: 机体质量 (单位：kg) 默认值0.78
        :param motor_assembly_weight: 单个电机组件质量 (单位：kg) 默认值0.055
        :return: 3x3惯性矩阵 (单位：kg·m²)
        """
        # 初始化3x3零矩阵
        inertia = np.zeros((3, 3), dtype=np.float32)

        # 计算机体惯性矩阵（均质长方体公式）
        x, y, z = body_box
        inertia[0, 0] = box_mass / 12 * (y ** 2 + z ** 2)  # Ixx
        inertia[1, 1] = box_mass / 12 * (x ** 2 + z ** 2)  # Iyy
        inertia[2, 2] = box_mass / 12 * (x ** 2 + y ** 2)  # Izz

        # 叠加旋翼惯性贡献（平行轴定理）
        for pos in rotor_poses:
            px, py, pz = pos
            inertia[0, 0] += motor_assembly_weight * (py ** 2 + pz ** 2)  # Ixx += m(y²+z²)
            inertia[1, 1] += motor_assembly_weight * (px ** 2 + pz ** 2)  # Iyy += m(x²+z²)
            inertia[2, 2] += motor_assembly_weight * (px ** 2 + py ** 2)  # Izz += m(x²+y²)

        return inertia

    def step(self, state, control):
        return self.A @ state + self.B @ control

    def get_dLTI(self):
        num_x = 12


        Ac = np.array([[0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
                       [0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
                       [0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
                       [0, 0, 0, 0, 0, 0, 0, -self.g, 0, 0, 0, 0],
                       [0, 0, 0, 0, 0, 0, self.g, 0, 0, 0, 0, 0],
                       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                       [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
                       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
                       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
                       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])
        Bc = np.array([[0, 0, 0, 0],
                       [0, 0, 0, 0],
                       [0, 0, 0, 0],
                       [0, 0, 0, 0],
                       [0, 0, 0, 0],
                       [1 / self.mass, 0, 0, 0],
                       [0, 0, 0, 0],
                       [0, 0, 0, 0],
                       [0, 0, 0, 0],
                       [0, 1/self.Ixx,0, 0],
                       [0, 0, 1/self.Iyy, 0],
                       [0, 0, 0, 1/self.Izz]])
        Cc = np.array([[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                       [0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                       [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                       [0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
                       [0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
                       [0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0]])
        Dc = np.zeros((6, 4))
        sysc = control.ss(Ac, Bc, Cc, Dc)
        # Discretization 双线性变换离散化
        sysd = control.sample_system(sysc, self.dt, method='bilinear')
        return sysd.A, sysd.B

class TrajectoryGenerator:
    def __init__(self, total_time):
        self.total_time = total_time

    def get_desired_state(self, t):
        # 参数设置
        radius = 7.0  # 水平螺旋半径（米）
        omega = 0.2  # 水平旋转角速度（rad/s）
        z_velocity = 0  # 垂直上升速度（米/秒）

        # 生成螺旋轨迹
        x_d = radius * np.cos(omega * t)
        y_d = radius * np.sin(omega * t)
        z_d = -3  # 垂直方向线性上升

        # 计算速度
        vx_d = -radius * omega * np.sin(omega * t)
        vy_d = radius * omega * np.cos(omega * t)
        vz_d = z_velocity  # 垂直速度恒定

        # 计算角度和角速度
        roll_d = 0.0  # 滚转角（保持水平）
        pitch_d = 0.0  # 俯仰角（保持水平）
        yaw_d = omega * t + np.pi / 2  # 航向角指向切线方向

        # 角速度
        # 角速度
        ro_d = 0.0  # 滚转角速度
        pi_d = 0.0  # 俯仰角速度
        ya_d = omega  # 偏航角速度（与水平旋转角速度一致）

        # 合并所有状态
        return np.array([x_d, y_d, z_d,
                         vx_d, vy_d, vz_d,
                         roll_d, pitch_d, yaw_d,
                         ro_d, pi_d, ya_d])


class MPCController:
    def __init__(self, quadrotor, N, Q, R,F_max,M_max):
        self.quadrotor = quadrotor
        self.N = N  # 预测时域
        self.Q = Q  # 状态权重矩阵（6x6）
        self.R = R  # 控制输入权重矩阵（4x4）
        self.F_max = F_max
        self.M_max = M_max

    def control(self, state, desired_states):
        n_states = self.quadrotor.n_states
        n_controls = self.quadrotor.n_controls
        x = cp.Variable((n_states, self.N + 1))
        u = cp.Variable((n_controls, self.N))

        cost = 0
        constraints = []

        # 初始状态约束
        constraints.append(x[:, 0] == state)

        for k in range(self.N):
            # 状态转移约束
            constraints.append(
                x[:, k + 1] ==
                self.quadrotor.A @ x[:, k] +
                self.quadrotor.B @ u[:, k]
            )

            # 代价函数
            cost += cp.quad_form(x[:, k] - desired_states[:, k], self.Q)
            cost += cp.quad_form(u[:, k], self.R)
            # 约束
            constraints.append(u[0, k] >= 0)  # 推力非负
            constraints.append(u[0, k] <= self.F_max)  # 最大推力
            constraints.append(u[1, k] >= -self.M_max)
            constraints.append(u[1, k] <= self.M_max)
            constraints.append(u[2, k] >= -self.M_max)
            constraints.append(u[2, k] <= self.M_max)
            constraints.append(u[3, k] >= -self.M_max)
            constraints.append(u[3, k] <= self.M_max)

        # 终端代价
        cost += cp.quad_form(x[:, -1] - desired_states[:, -1], self.Q)

        prob = cp.Problem(cp.Minimize(cost), constraints)
        prob.solve(solver=cp.OSQP)

        return u[:, 0].value



def get_drone_state(client, buffer):
    state = client.getMultirotorState()
    kinematics = state.kinematics_estimated

    # 直接写入预分配的缓冲区
    buffer[0:3] = [kinematics.position.x_val,
                   kinematics.position.y_val,
                   kinematics.position.z_val]

    buffer[3:6] = [kinematics.linear_velocity.x_val,
                   kinematics.linear_velocity.y_val,
                   kinematics.linear_velocity.z_val]

    angles = airsim.to_eularian_angles(kinematics.orientation)

    buffer[6] = angles[1]  # roll
    buffer[7] = angles[0]  # pitch
    buffer[8] = angles[2]  # yaw

    buffer[9:12] = [kinematics.angular_velocity.x_val,
                    kinematics.angular_velocity.y_val,
                    kinematics.angular_velocity.z_val]
    return buffer

def fM2u(f, M):
    mat = np.array([[4.179446268,       4.179446268,        4.179446268,        4.179446268],
                    [-0.6723341164784,  0.6723341164784,    0.6723341164784,    -0.6723341164784],
                    [0.6723341164784,   -0.6723341164784,   0.6723341164784,    -0.6723341164784],
                    [0.055562,          0.055562,           -0.055562,          -0.055562]])
    fM = np.vstack([f, M])
    u = np.dot(np.linalg.inv(mat), fM)
    u1 = u[0, 0]
    u2 = u[1, 0]
    u3 = u[2, 0]
    u4 = u[3, 0]
    return u1, u2, u3, u4


def simulate():
    total_time = 40

    quadrotor = Quadrotor(dt=0.2)
    traj_gen = TrajectoryGenerator(total_time=total_time)
    N = 8 # 预测时域
    Q = np.diag([1.0, 1.0, 1.0, 0.1, 0.1, 0.1,1.0, 1.0, 1.0, 0.1, 0.1, 0.1])  # 状态权重（位置权重高）
    R = np.diag([0.1, 0.01, 0.01, 0.01])  # 输入权重（总力权重较高）
    F_max = 4 * 4.179  # 4个电机最大总推力（假设每个电机最大推力4.179对应u=1）
    M_max = 0.055562  # 根据电机力矩能力调整
    mpc = MPCController(quadrotor, N, Q, R, F_max, M_max)

    num_steps = int(total_time / quadrotor.dt)
    time_steps = np.linspace(0, total_time, num_steps)
    actual_trajectory = []
    desired_trajectory = []
    p_list = []
    for i in range(num_steps):

        t = i * quadrotor.dt
        desired_state = traj_gen.get_desired_state(t)
        desired_trajectory.append(desired_state)
        p_list.append(airsim.Vector3r(desired_state[0], desired_state[1], desired_state[2]))
    desired_trajectory = np.array(desired_trajectory)
    client = airsim.MultirotorClient()
    client.enableApiControl(True)
    client.takeoffAsync().join()
    client.simPlotLineStrip(p_list, color_rgba=[0, 1, 0, 0], is_persistent=True)
    client.simSetTraceLine(color_rgba=[1, 0, 0], thickness=5)
    for i in range(num_steps-N):
        state_buffer = np.empty(12)
        state=get_drone_state(client, state_buffer)

        # 提取当前和未来 N 步的期望状态
        desired_states = desired_trajectory[i:i + N + 1].T  # 形状：(12, N+1)

        # 计算控制输入
        control = mpc.control(state, desired_states)
        actual_trajectory.append(state[:3])
        # 更新状态
        f=control[0]
        M=np.array([[control[1]],[control[2]],[control[3]]])
        u1, u2, u3, u4 = fM2u(f, M)
        print("u1",u1)
        print("u2",u2)
        print("u3",u3)
        print("u4",u4)
        client.moveByMotorPWMsAsync(u1, u2, u3, u4, quadrotor.dt)
        time.sleep(quadrotor.dt)
        # 存储轨迹


    #绘制三维轨迹
    # fig = plt.figure(figsize=(12, 6))
    # ax = fig.add_subplot(111, projection='3d')
    # actual = np.array(actual_trajectory)
    # desired = np.array(desired_trajectory)
    #
    # ax.plot(desired[:, 0], desired[:, 1], desired[:, 2], 'b--', label='Desired')
    # ax.plot(actual[:, 0], actual[:, 1], actual[:, 2], 'r-', label='Actual')
    # ax.set_xlabel('X (m)'), ax.set_ylabel('Y (m)'), ax.set_zlabel('Z (m)')
    # ax.legend()
    # ax.set_title('3D Trajectory Tracking')
    # plt.figure()
    # plt.plot(time_steps, actual[:,0], 'r-', label='Actual')
    # plt.plot(time_steps, desired[:,0], 'b--', label='Desired')
    # plt.xlabel('time (s)'), plt.ylabel('X (m)')
    # plt.legend()
    # plt.title('X Trajectory Tracking')
    # plt.figure()
    # plt.plot(time_steps, actual[:, 1], 'r-', label='Actual')
    # plt.plot(time_steps, desired[:, 1], 'b--', label='Desired')
    # plt.xlabel('time (s)'), plt.ylabel('Y (m)')
    # plt.legend()
    # plt.title('Y Trajectory Tracking')
    # plt.figure()
    # plt.plot(time_steps, actual[:, 2], 'r-', label='Actual')
    # plt.plot(time_steps, desired[:, 2], 'b--', label='Desired')
    # plt.xlabel('time (s)'), plt.ylabel('Z (m)')
    # plt.legend()
    # plt.title('Z Trajectory Tracking')
    # plt.tight_layout()
    # plt.show()


if __name__ == "__main__":
    simulate()
